##########################################################################################

library(data.table)
library(optparse)
library(Seurat)
library(ArchR)
library(ggsci)
library(tidyverse)

##########################################################################################
option_list <- list(
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--rna_file"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 整合atac和rna的文件
    comine_data_file <- "~/20231121_singleMuti/results/qc_atac_v2/all/testis_combined_peak.combineRNA.qc.Rdata"

    ## 单细胞表达文件
    rna_file <- "~/20231121_singleMuti/results/qc_atac_v2/all/testis_combined.annotationCellType.qc.Rdata"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/qc_atac_v3"

}


###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

comine_data_file <- opt$comine_data_file
rna_file <- opt$rna_file
out_path <- opt$out_path

dir.create(out_path , recursive = T)
dir.create(paste0(out_path , "/all") , recursive = T)
dir.create(paste0(out_path , "/germ") , recursive = T)
dir.create(paste0(out_path , "/somatic") , recursive = T)

###########################################################################################

a <- load(comine_data_file)
## testis_combined_peak_combineRNA

b <- load(rna_file)
## scrnat

###########################################################################################
## 细胞颜色
use_colors <- c(pal_npg("nrc")(10) , pal_jco("default")(6))
#names(use_colors) <- unique(scrnat$cell_type)
names(use_colors) <-c( "Myoid cells","Leydig cells" ,          
"Endothelial cells","Zygotene",         
"Round&ElongateS.tids","Patchytene",
"SSC","Sperm" ,                 
"Diplotene","Early stage of spermatids", 
"Leptotene","Sertoli cells",            
"Macrophages","Differenting&Differented SPG",
"Pericytes","NKT cells")

###########################################################################################
## 生殖细胞
cell_level_germ <- c("SSC","Differenting&Differented SPG", "Leptotene",
    "Zygotene","Patchytene","Diplotene",
    "Early stage of spermatids","Round&ElongateS.tids","Sperm"
    )

## 体细胞
cell_level_somatic <- c( "Myoid cells","Leydig cells" , "Endothelial cells", 
"Sertoli cells" , "Macrophages" , 
"Pericytes" , "NKT cells")

###########################################################################################
## 对于RNA，看cluster里面哪一簇和之前不一致
scrnat$seurat_clusters_raw <- scrnat$seurat_clusters

## 重新聚类
scrnat <- FindNeighbors(scrnat, reduction = "pca", dims = 1:30)
scrnat <- FindClusters(scrnat , graph.name = grep( "snn" , names(scrnat@graphs) , value = T ) , resolution = 0.6)

out_file <- paste0( out_path , "/magic3_cluster.qc.cell_type.pdf" ) 
p <- DimPlot(scrnat, reduction = "umap", label = TRUE, group.by = 'cell_type' , cols = use_colors ) + theme(legend.position = 'none') 
ggsave( out_file , p ,width = 6 , height = 6)
out_file <- paste0( out_path , "/magic3_cluster.qc.seurat_clusters.pdf" ) 
p <- DimPlot(scrnat, reduction = "umap", label = TRUE, group.by = 'seurat_clusters') + theme(legend.position = 'none') 
ggsave( out_file , p ,width = 6 , height = 6)

## 每个细胞类型在那个类里面
df_cellClus <- table(scrnat$seurat_clusters , scrnat$cell_type)

## 每个细胞类型均有一个主要的cluster，去除少量的cluster
all_dell_cell <- c()
for( cellT in c("Zygotene" , "Myoid cells" , "Endothelial cells") ){
  print(cellT)

  ## 准备去除的细胞，该细胞类型所在的cluster小于50个细胞在里面
  del_clus <- names(which(df_cellClus[,cellT] < 50 & df_cellClus[,cellT] > 0 ))

  if(length(del_clus) > 0){
    ## 准备去除的细胞
    del_cell <- subset(scrnat ,  (seurat_clusters %in% del_clus) & cell_type == cellT )$cell
    all_dell_cell <- c( all_dell_cell , names((del_cell)) )
  }
}

tmp_cell <- scrnat@meta.data$cell[!scrnat@meta.data$cell %in% all_dell_cell]
scrnat_tmp <- subset( scrnat , (cell %in% tmp_cell) )
#scrnat_tmp <- RunUMAP(scrnat_tmp, dims = 1:10)
## 可视化
out_file <- paste0( out_path , "/magic3_cluster.qc.cell_type.tmp.pdf" ) 
p <- DimPlot(scrnat_tmp, reduction = "umap", label = TRUE, group.by = 'cell_type' , cols = use_colors ) + theme(legend.position = 'none') 
ggsave( out_file , p ,width = 6 , height = 6)
out_file <- paste0( out_path , "/magic3_cluster.qc.seurat_clusters.tmp.pdf" ) 
p <- DimPlot(scrnat_tmp, reduction = "umap", label = TRUE, group.by = 'seurat_clusters') + theme(legend.position = 'none') 
ggsave( out_file , p ,width = 6 , height = 6)


## kmeans聚类更优秀
## 去除一些散的点
umap_type = scrnat_tmp@reductions$umap@cell.embeddings %>%
      as.data.frame() %>% cbind(type = scrnat_tmp@meta.data$cell_type)

km <- kmeans(umap_type[,1:2], 20, nstart = 20)
km <- data.frame( cell = scrnat_tmp@meta.data$cell , cell_type = scrnat_tmp@meta.data$cell_type , cluster = km$cluster)
clus_km <- table(km$cluster , km$cell_type)

all_dell_cell_km <- c()
for( cellT in c("Myoid cells" , "Zygotene" , "Endothelial cells" , "Pericytes" , "Leydig cells") ){
  print(cellT)

  ## 准备去除的细胞，该细胞类型所在的cluster小于25个细胞在里面
  del_clus <- names(which(clus_km[,cellT] < 30 & clus_km[,cellT] > 0 ))
  if(length(del_clus) > 0){
    ## 准备去除的细胞
    del_cell <- subset(km ,  (cluster %in% del_clus) & cell_type == cellT )$cell
    all_dell_cell_km <- c( all_dell_cell_km , del_cell )
  }
}

## 和之前聚类的去除的细胞取并集
all_dell_cell <- unique(c(all_dell_cell , all_dell_cell_km))
tmp_cell <- scrnat@meta.data$cell[!scrnat@meta.data$cell %in% all_dell_cell]
scrnat_tmp <- subset( scrnat_tmp , (cell %in% tmp_cell) )
#scrnat_tmp <- RunUMAP(scrnat_tmp, dims = 1:10)

## 可视化
out_file <- paste0( out_path , "/magic3_cluster.qc.cell_type.tmp.pdf" ) 
p <- DimPlot(scrnat_tmp, reduction = "umap", label = TRUE, group.by = 'cell_type' , cols = use_colors ) + theme(legend.position = 'none') 
ggsave( out_file , p ,width = 6 , height = 6)
out_file <- paste0( out_path , "/magic3_cluster.qc.seurat_clusters.tmp.pdf" ) 
p <- DimPlot(scrnat_tmp, reduction = "umap", label = TRUE, group.by = 'seurat_clusters') + theme(legend.position = 'none') 
ggsave( out_file , p ,width = 6 , height = 6)


###########################################################################################
## ATAC和RNA的取并集
all_dell_cell_use <- gsub( "_" , "#" ,  all_dell_cell )
## 去除185个细胞

## 最后使用的细胞
use_cell_all <- rownames(testis_combined_peak_combineRNA@cellColData)[!rownames(testis_combined_peak_combineRNA@cellColData) %in% all_dell_cell_use]

## 生殖细胞
subCells_germ <- getCellNames(testis_combined_peak_combineRNA)[as.character(testis_combined_peak_combineRNA@cellColData[["cell_type"]]) %in% cell_level_germ]
## 最后使用
use_cell_germ <- subCells_germ[subCells_germ %in% use_cell_all]

## 体细胞
subCells_somatic <- getCellNames(testis_combined_peak_combineRNA)[as.character(testis_combined_peak_combineRNA@cellColData[["cell_type"]]) %in% cell_level_somatic]
## 最后使用
use_cell_somatic <- subCells_somatic[subCells_somatic %in% use_cell_all]

###########################################################################################

scrnat$seurat_clusters <- scrnat$seurat_clusters_raw

###########################################################################################
## germ
## 输出到对应目录
projHeme5 <- subsetArchRProject(
  ArchRProj = testis_combined_peak_combineRNA,
  cells = use_cell_germ,
  outputDirectory = paste0(out_path , "/germ"),
  dropCells = TRUE,
  force = TRUE
)

projHeme5 <- addUMAP( projHeme5 , seed = 1 , force = TRUE)
p <- plotEmbedding(ArchRProj = projHeme5, colorBy = "cellColData", name = "cell_type", embedding = "UMAP" , pal = use_colors)
p1 <- plotEmbedding(ArchRProj = projHeme5, colorBy = "cellColData", name = "Clusters", embedding = "UMAP")
image_name <- paste0( out_path , "/germ/plotEmbedding.atac.pdf")
pdf(image_name , width = 10, height = 8)
print(p)
print(p1)
dev.off()

## 生殖细胞最后保留的细胞
projHeme5_germ <- projHeme5
scrnat_germ <- subset( scrnat , cell %in% gsub( "#" , "_" , use_cell_germ ) )

###########################################################################################
## somatic
## 输出到对应目录
projHeme5 <- subsetArchRProject(
  ArchRProj = testis_combined_peak_combineRNA,
  cells = use_cell_somatic,
  outputDirectory = paste0(out_path , "/somatic"),
  dropCells = TRUE,
  force = TRUE
)

projHeme5 <- addUMAP( projHeme5 , seed = 1 , force = TRUE)
p <- plotEmbedding(ArchRProj = projHeme5, colorBy = "cellColData", name = "cell_type", embedding = "UMAP" , pal = use_colors)
p1 <- plotEmbedding(ArchRProj = projHeme5, colorBy = "cellColData", name = "Clusters", embedding = "UMAP")
image_name <- paste0( out_path , "/somatic/plotEmbedding.atac.pdf")
pdf(image_name , width = 10, height = 8)
print(p)
print(p1)
dev.off()

## 生殖细胞最后保留的细胞
projHeme5_somatic <- projHeme5
scrnat_somatic <- subset( scrnat , cell %in% gsub( "#" , "_" , use_cell_somatic ) )

###########################################################################################
## 所有的
subCells <- c(use_cell_germ , use_cell_somatic)

projHeme5 <- subsetArchRProject(
  ArchRProj = testis_combined_peak_combineRNA,
  cells = subCells,
  outputDirectory = paste0(out_path , "/all"),
  dropCells = TRUE,
  force = TRUE
)

## 重聚类
projHeme5 <- addUMAP( projHeme5 , seed = 1 , force = TRUE)

## atac的聚类图,标记细胞类型
p <- plotEmbedding(ArchRProj = projHeme5, colorBy = "cellColData", name = "cell_type", embedding = "UMAP" , pal = use_colors)
p1 <- plotEmbedding(ArchRProj = projHeme5, colorBy = "cellColData", name = "Clusters", embedding = "UMAP")
image_name <- paste0( out_path , "/all/plotEmbedding.atac.pdf")
pdf(image_name , width = 10, height = 8)
print(p)
print(p1)
dev.off()

## 所有最后用到的细胞
projHeme5_all <- projHeme5
scrnat_all <- subset( scrnat , cell %in% gsub( "#" , "_" , subCells ) )

###########################################################################################
## 输出
## 胚系的
scrnat <- scrnat_germ
testis_combined_peak_combineRNA <- projHeme5_germ
image_name <- paste0( out_path , "/germ/testis_combined_peak.combineRNA.qc.Rdata")
save(  testis_combined_peak_combineRNA , file = image_name )
saveArchRProject( testis_combined_peak_combineRNA , paste0( out_path , "/germ/") )
image_name <- paste0( out_path , "/germ/testis_combined.annotationCellType.qc.Rdata")
save(  scrnat , file = image_name )

## 体细胞的
scrnat <- scrnat_somatic
testis_combined_peak_combineRNA <- projHeme5_somatic
image_name <- paste0( out_path , "/somatic/testis_combined_peak.combineRNA.qc.Rdata")
save(  testis_combined_peak_combineRNA , file = image_name )
saveArchRProject( testis_combined_peak_combineRNA , paste0( out_path , "/somatic/") )
image_name <- paste0( out_path , "/somatic/testis_combined.annotationCellType.qc.Rdata")
save(  scrnat , file = image_name )

## 所有的
scrnat <- scrnat_all
testis_combined_peak_combineRNA <- projHeme5_all
image_name <- paste0( out_path , "/all/testis_combined_peak.combineRNA.qc.Rdata")
save(  testis_combined_peak_combineRNA , file = image_name )
saveArchRProject( testis_combined_peak_combineRNA , paste0( out_path , "/all/") )
image_name <- paste0( out_path , "/all/testis_combined.annotationCellType.qc.Rdata")
save(  scrnat , file = image_name )
